import tensorflow as tf  # 导入TensorFlow库
from keras.preprocessing.image import ImageDataGenerator
from keras import layers, Model, Sequential, activations, optimizers, losses  # 从Keras模块导入需要的层、模型、优化器和损失函数
from keras.datasets import cifar10  # 导入CIFAR-10数据集

# 加载CIFAR-10数据集并进行数据预处理
(x_train, y_train), (x_test, y_test) = cifar10.load_data()  # 加载训练和测试数据
x_train = x_train.astype('float32') / 255.0  # 将训练数据标准化到0-1范围
x_test = x_test.astype('float32') / 255.0  # 将测试数据标准化到0-1范围

# 数据增强
data_augmentation = ImageDataGenerator(
    width_shift_range=0.1,  # 随机水平移动
    height_shift_range=0.1,  # 随机垂直移动
    horizontal_flip=True,  # 随机水平翻转
)
# 定义AlexNet模型
class AlexNet(Model):  # 继承Model类定义AlexNet模型
    def __init__(self):  # 初始化模型
        super(AlexNet, self).__init__()  # 调用父类的初始化方法
        self.conv = Sequential([  # 定义卷积层
            layers.Conv2D(filters=6, kernel_size=(3, 3), activation='relu', padding='same'),  # 第一层卷积层
            layers.MaxPooling2D(),  # 第一层池化层
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same'),  # 第二层卷积层
            layers.MaxPooling2D(),  # 第二层池化层
            layers.Conv2D(filters=24, kernel_size=(3, 3), activation='relu', padding='same'),  # 第三层卷积层
            layers.Conv2D(filters=24, kernel_size=(3, 3), activation='relu', padding='same'),  # 第四层卷积层
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same'),  # 第五层卷积层
            layers.MaxPooling2D()  # 第三层池化层
        ])
        self.flat = layers.Flatten()  # 定义展平层
        self.fc = Sequential([  # 定义全连接层
            layers.Dense(128, activation='relu'),  # 第一层全连接层
            layers.Dense(64, activation='relu'),  # 第一层全连接层
            layers.Dense(64, activation='relu'),  # 第二层全连接层
            layers.Dense(10, activation='softmax')  # 输出层，10个类别，使用softmax激活
        ])

    def call(self, inputs):  # 定义前向传播
        x = self.conv(inputs)  # 输入数据经过卷积层
        x = self.flat(x)  # 展平卷积层输出
        x = self.fc(x)  # 经过全连接层
        return x  # 返回输出结果


model = AlexNet()  # 创建AlexNet模型实例
model.build(input_shape=(None, 32, 32, 3))  # 定义输入形状
model.summary()
model.compile(optimizer=optimizers.Adam(),  # 编译模型，使用Adam优化器
              loss=losses.sparse_categorical_crossentropy,  # 使用稀疏分类交叉熵损失函数
              metrics=['accuracy'])  # 评价指标为准确率

model.summary()  # 打印模型概要

model.fit(data_augmentation.flow(x_train, y_train, batch_size=64), epochs=10)  # 训练模型，设置批次大小和训练轮数

test_loss, test_acc = model.evaluate(x_test, y_test)  # 在测试集上评估模型
print(f'Test accuracy: {test_acc}')  # 打印测试准确率

model.save_weights('alexnet_cifar10.h5')  # 保存模型权重
